# -*- coding: utf-8 -*
from flyai.framework import FlyAI
from net import get_model
from torchvision import transforms
import torch
import os
from path import MODEL_PATH
from PIL import Image
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.checkpoint import DetectionCheckpointer
from detectron2.data import MetadataCatalog, DatasetCatalog
from detectron2.utils.visualizer import ColorMode
from detectron2.utils.visualizer import Visualizer
import cv2
from detectron2 import model_zoo

# 判断gpu是否可用
if torch.cuda.is_available():
    device = torch.device('cuda')
else:
    device = torch.device('cpu')

threshold = 0.2

def get_return(boxes, labels, scores, image_name, label_id_name):
    ''' 输入示例：
    :param boxes: [[735.097,923.59283,770.1911,998.49335], [525.39496,535.89667,578.4822,589.8431 ]]]   box格式[Xmin, Ymin, Xmax, Ymax]
    :param labels: [1,1]
    :param scores: [0.2， 0.3]
    :param image_name: 0.jpg
    :param label_id_name: {1: 'TBbacillus'}  label id 到 label name的映射

    :return:  [{"image_name": '0.jpg', "label_name": 'TBbacillus', "bbox": [735, 923, 770-735, 998-923], "confidence": 0.2},
                {"image_name": '0.jpg', "label_name": 'TBbacillus', "bbox": [525, 535, 578-525, 589-535], "confidence": 0.3}]
                返回 box的格式为[xmin, ymin, width, height]
    最终评估方式采用coco数据集的map，具体可参考 https://cocodataset.org/
    '''
    result = []
    if len(boxes) == len(labels) == len(scores):
        for i in range(len(boxes)):
            box = boxes[i] # [Xmin, Ymin, Xmax, Ymax]
            bbox = [int(box[0]), int(box[1]), int(box[2]-box[0]), int(box[3]-box[1])] # [xmin, ymin, width, height]
            #label_name = label_id_name[labels[i]]
            label_name = 'TBbacillus'
            confidence = scores[i]
            ann = {"image_name": image_name, "label_name": label_name, "bbox": bbox, "confidence": confidence}
            result.append(ann)
    return result



class Prediction(FlyAI):
    def cv2_imshow(self,a,boxes,scores,classes):
        """A replacement for cv2.imshow() for use in Jupyter notebooks.
        
         Args:
           a : np.ndarray. shape (N, M) or (N, M, 1) is an NxM grayscale image. shape
             (N, M, 3) is an NxM BGR color image. shape (N, M, 4) is an NxM BGRA color
             image.
        """
        a = a.clip(0, 255).astype('uint8')
        # cv2 stores colors as BGR; convert to RGB
        if a.ndim == 3:
          if a.shape[2] == 4:
            a = cv2.cvtColor(a, cv2.COLOR_BGRA2RGBA)
          else:
            a = cv2.cvtColor(a, cv2.COLOR_BGR2RGB)
        for box,score,cla in zip(boxes,scores,classes):
            xmin,ymin,Xmax,Ymax = box
            cv2.rectangle(a,(xmin,ymin),(Xmax,Ymax),(0,0,255),3)
            cv2.putText(a,"{:.2f} {}".format(score,cla),(xmin,ymin),cv2.FONT_HERSHEY_COMPLEX_SMALL,0.8,(255,0,0))

        #cv2.imwrite('./save/show_ing.jpg',a)
        cv2.imshow("display",a)
        cv2.waitKey(0)
        cv2.destroyAllWindows()

    def load_model(self):
        '''
        模型初始化，必须在此方法中加载模型
        '''
        cfg = get_cfg()
        #cfg.merge_from_file(model_zoo.get_config_file("COCO-Detection/retinanet_R_50_FPN_1x.yaml"))
        #cfg.merge_from_file(model_zoo.get_config_file("COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x.yaml"))
        cfg.merge_from_file(model_zoo.get_config_file("COCO-Detection/faster_rcnn_R_101_FPN_3x.yaml"))
        cfg.MODEL.WEIGHTS = os.path.join(MODEL_PATH, "model_final.pth")
        cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7   # set a custom testing threshold for this model
        cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128
        cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1 
        self.predictor = DefaultPredictor(cfg)
        self.TBDetection_metadata = MetadataCatalog.get("TBDetection_train")
        print('load model done...')

    def predict(self, image_path):
        '''
        模型预测返回结果
        :param input: 评估传入样例 {"image_path": "./data/input/image/0.jpg"}
        :return: 具体见 get_return 方法
        '''
        img = cv2.imread(image_path)
        image_name = image_path.split('/')[-1]

        outputs = self.predictor(img)

        outputs = outputs["instances"]
        boxes = outputs.pred_boxes.tensor.to('cpu')
        scores = outputs.scores.to("cpu")
        classes = outputs.pred_classes.to("cpu")
        tmp = (scores > threshold).nonzero().to('cpu').squeeze().numpy().tolist()

        selectScores = torch.index_select(scores,0,torch.LongTensor(tmp)).numpy()
        selectBoxs = torch.index_select(boxes,0,torch.LongTensor(tmp)).numpy()
        selectClass = torch.index_select(classes,0,torch.LongTensor(tmp)).numpy()


        #self.cv2_imshow(img,selectBoxs,selectScores,selectClass)
        result = get_return(selectBoxs, selectClass, selectScores, image_name, label_id_name={0: 'TBbacillus'})
        print(result)
        return result

if __name__ == "__main__":
    pre = Prediction()
    pre.load_model()
    result = pre.predict(image_path="./data/input/TBDetection/image/1811.jpg")
    print(result)